In [906]:
import sys
# Third-party
import astropy.coordinates as coord
import astropy.units as u
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
plt.style.use('apw-notebook')
%matplotlib inline
from scipy.special import gamma as Gamma
from scipy.stats import gamma, norm
from scipy.misc import logsumexp
In [907]:
n = 3
s = 15.
x1 = norm.rvs([0]*n, scale=[s]*n, size=(100000,n))
x2 = norm.rvs([0]*n, scale=[s]*n, size=(100000,n))
d12 = np.linalg.norm(x1 - x2, axis=1)
In [908]:
def pdf(x, s, n):
V = (2*s**2)
A = 2**(n/2-1) * V**(n/2) * Gamma(n/2)
return x**(n-1) * np.exp(-x**2/(2*V)) / A
In [910]:
plt.hist(d12, bins='auto', normed=True);
_xs = np.linspace(0.1, 100, 256)
plt.plot(_xs, pdf(_xs, s, n))
Out[910]:
In [921]:
from scipy.stats import scoreatpercentile
In [941]:
cov = np.eye(3) * np.array([25., 15., 12.])**2 # Nordstrom 2004
# cov = np.eye(3) * np.array([30., 18., 13.])**2 # Sharma 2014
x1 = np.random.multivariate_normal([0,0,0.], cov, size=100000)
x2 = np.random.multivariate_normal([0,0,0.], cov, size=100000)
d12 = np.linalg.norm(x1 - x2, axis=1)
plt.hist(d12, bins='auto', normed=True);
_xs = np.linspace(0.1, 100, 256)
# plt.plot(_xs, pdf(_xs, np.sqrt(np.trace(cov)/3), n))
plt.plot(_xs, pdf(_xs, 17, n))
print(np.sqrt(np.trace(cov)/3))
plt.axvline(np.sqrt(np.trace(cov)), color='tab:red')
for p in scoreatpercentile(d12, [15, 85]):
print(p)
plt.axvline(p, color='tab:red', alpha=0.6, linestyle='dashed')
# plt.text(p-4)
In [187]:
n_data = 400
qu_n_data = n_data//4
xyz1 = coord.CartesianRepresentation(100 * np.random.random(size=(3, n_data))*u.pc)
dx = coord.PhysicsSphericalRepresentation(r=np.random.random(n_data)*u.pc,
phi=np.random.uniform(0,2*np.pi,n_data)*u.rad,
theta=np.arccos(2*np.random.random(n_data) - 1)*u.rad)
xyz2 = xyz1 + dx
tmp = coord.PhysicsSphericalRepresentation(r=np.ones(n_data),
phi=np.random.uniform(0,2*np.pi,n_data)*u.rad,
theta=np.arccos(2*np.random.random(n_data) - 1)*u.rad)
tmp = tmp.represent_as(coord.CartesianRepresentation)
vxyz1 = np.random.normal(0, 30, n_data)*u.km/u.s * tmp.represent_as(coord.CartesianDifferential)
vxyz2 = vxyz1[np.concatenate((np.arange(qu_n_data),
qu_n_data+np.random.choice(3*qu_n_data, size=3*qu_n_data, replace=False)))]
In [189]:
plt.hist((vxyz1 - vxyz2).norm(), bins='auto');
In [190]:
icrs1 = coord.ICRS(xyz1.with_differentials(vxyz1))
icrs2 = coord.ICRS(xyz2.with_differentials(vxyz2))
In [717]:
pm_err1 = np.random.uniform(0.2, 1., size=(n_data, 2)) # mas/yr
pm_err2 = np.random.uniform(0.2, 1., size=(n_data, 2))
rv_err1 = np.random.uniform(0.5, 1., size=n_data) # km/s
rv_err2 = np.random.uniform(0.5, 1., size=n_data)
pm1 = np.random.normal(np.stack((icrs1.pm_ra_cosdec.to(u.mas/u.yr).value,
icrs1.pm_dec.to(u.mas/u.yr).value)).T,
pm_err1)
pm2 = np.random.normal(np.stack((icrs2.pm_ra_cosdec.to(u.mas/u.yr).value,
icrs2.pm_dec.to(u.mas/u.yr).value)).T,
pm_err2)
rv1 = np.random.normal(icrs1.radial_velocity.to(u.km/u.s).value,
rv_err1)
rv2 = np.random.normal(icrs2.radial_velocity.to(u.km/u.s).value,
rv_err2)
In [718]:
def get_tangent_basis(ra, dec):
"""
row vectors are the tangent-space basis at (alpha, delta, r)
ra, dec in radians
"""
ra = np.atleast_1d(ra)
dec = np.atleast_1d(dec)
s = []
for cra, cdec in zip(ra, dec):
s.append(
np.array([
[-np.sin(cra), np.cos(cra), 0.],
[-np.sin(cdec)*np.cos(cra), -np.sin(cdec)*np.sin(cra), np.cos(cdec)],
[np.cos(cdec)*np.cos(cra), np.cos(cdec)*np.sin(cra), np.sin(cdec)]
]).T)
return np.array(s).squeeze()
In [719]:
fac = 1 / (1*u.km/u.s).to(u.pc*u.mas/u.yr, u.dimensionless_angles()).value
In [720]:
M1 = get_tangent_basis(icrs1.ra.radian, icrs1.dec.radian)
M2 = get_tangent_basis(icrs2.ra.radian, icrs2.dec.radian)
From tests, it seems like the log-likelihood values converge at ~512 samples or so
In [721]:
n_samples = 512
pm1_samples = np.random.normal(pm1, pm_err1, size=(n_samples, n_data, 2))
pm2_samples = np.random.normal(pm2, pm_err2, size=(n_samples, n_data, 2))
rv1_samples = np.random.normal(rv1, rv_err1, size=(n_samples, n_data))
rv2_samples = np.random.normal(rv2, rv_err2, size=(n_samples, n_data))
v1_samples = np.array([icrs1.distance.value[None] * pm1_samples[...,0] * fac,
icrs1.distance.value[None] * pm1_samples[...,1] * fac,
rv1_samples])
v2_samples = np.array([icrs2.distance.value[None] * pm2_samples[...,0] * fac,
icrs2.distance.value[None] * pm2_samples[...,1] * fac,
rv2_samples])
v1_samples = np.einsum('nij,ikn->jkn', M1, v1_samples)
v2_samples = np.einsum('nij,ikn->jkn', M2, v2_samples)
dv_samples = np.linalg.norm(v1_samples - v2_samples, axis=0)
dv_samples.shape
Out[721]:
In [722]:
def ln_dv_pdf(x, sigma):
return 2*np.log(x) - x**2/(4*sigma**2) - 1.2655121234846456 - 3*np.log(sigma)
def ln_likelihood(p):
f = p[0]
term1 = ln_dv_pdf(dv_samples, 0.1)
term2 = ln_dv_pdf(dv_samples, 30.)
ll = np.zeros(term1.shape[1])
for n in range(n_data):
ll[n] += logsumexp([term1[:,n]+log(f), term2[:,n]+log(1-f)]) - np.log(n_samples)
return ll
In [483]:
%timeit ln_likelihood([0.5])
In [462]:
lls = []
fs = np.linspace(0, 1, 128)
for f in fs:
lls.append(ln_likelihood([f]).sum())
lls = np.array(lls)
plt.plot(fs, np.exp(lls - lls.max()))
Out[462]:
In [832]:
v1_samples_tmp = np.array([pm1_samples[...,0] * fac,
pm1_samples[...,1] * fac,
rv1_samples])
v2_samples_tmp = np.array([pm2_samples[...,0] * fac,
pm2_samples[...,1] * fac,
rv2_samples])
def get_dv_samples1(d1, d2):
v1_tmp = v1_samples_tmp * np.vstack((d1, d1, np.ones_like(d1)))[:,None]
v2_tmp = v2_samples_tmp * np.vstack((d2, d2, np.ones_like(d2)))[:,None]
v1_samples = np.einsum('nji,ikn->jkn', M1, v1_tmp)
v2_samples = np.einsum('nji,ikn->jkn', M2, v2_tmp)
return np.linalg.norm(v1_samples - v2_samples, axis=0)
def get_dv_samples2(d1, d2):
v1_tmp = v1_samples_tmp * np.vstack((d1, d1, np.ones_like(d1)))[:,None]
v2_tmp = v2_samples_tmp * np.vstack((d2, d2, np.ones_like(d2)))[:,None]
v1_samples = np.array([M1[n].dot(v1_tmp[...,n]) for n in range(n_data)])
v2_samples = np.array([M2[n].dot(v2_tmp[...,n]) for n in range(n_data)])
return np.linalg.norm(v1_samples - v2_samples, axis=1).T
def get_dv_samples3(d1, d2):
v1_tmp = (v1_samples_tmp * np.vstack((d1, d1, np.ones_like(d1)))[:,None]).T
v2_tmp = (v2_samples_tmp * np.vstack((d2, d2, np.ones_like(d2)))[:,None]).T
v1_samples = (M1[:,None] * v1_tmp[:,:,None]).sum(axis=-1)
v2_samples = (M2[:,None] * v2_tmp[:,:,None]).sum(axis=-1)
return np.linalg.norm(v1_samples - v2_samples, axis=-1).T
In [833]:
get_dv_samples = get_dv_samples2
In [813]:
plx1 = icrs1.distance.to(u.mas, u.parallax()).value
plx2 = icrs1.distance.to(u.mas, u.parallax()).value
plx1_err = np.full(n_data, 0.3)
plx2_err = np.full(n_data, 0.3)
plx1 = np.random.normal(plx1, plx1_err)
plx2 = np.random.normal(plx2, plx2_err)
In [814]:
n_dgrid = 5
d1_min, d1_max = 1000 / (plx1 + 3*plx1_err), 1000 / (plx1 - 3*plx1_err)
d2_min, d2_max = 1000 / (plx2 + 3*plx2_err), 1000 / (plx2 - 3*plx2_err)
d1_grids = np.array([np.linspace(d1_min[i], d1_max[i], n_dgrid) for i in range(n_data)])
d2_grids = np.array([np.linspace(d2_min[i], d2_max[i], n_dgrid) for i in range(n_data)])
In [827]:
assert np.allclose(get_dv_samples1(d1_grids[:,2], d2_grids[:,2]),
get_dv_samples2(d1_grids[:,2], d2_grids[:,2]))
assert np.allclose(get_dv_samples1(d1_grids[:,2], d2_grids[:,2]),
get_dv_samples3(d1_grids[:,2], d2_grids[:,2]))
In [736]:
%timeit get_dv_samples1(d1_grids[:,2], d2_grids[:,2])
%timeit get_dv_samples2(d1_grids[:,2], d2_grids[:,2])
%timeit get_dv_samples3(d1_grids[:,2], d2_grids[:,2])
In [324]:
from scipy.integrate import simps
In [839]:
def ln_likelihood_at_d1d2(p, d1, d2):
f = p[0]
# b = np.vstack((np.full(n_samples, f), np.full(n_samples, 1-f)))
dv_samples = get_dv_samples(d1, d2)
term1 = ln_dv_pdf(dv_samples, 1.) + log(f)
term2 = ln_dv_pdf(dv_samples, 25.) + log(1-f)
return (logsumexp(term1, axis=0) - log(n_samples),
logsumexp(term2, axis=0) - log(n_samples))
def ln_likelihood_with_dist(p):
ll_grid1 = np.zeros((n_data, n_dgrid, n_dgrid))
ll_grid2 = np.zeros((n_data, n_dgrid, n_dgrid))
terms = ln_likelihood_at_d1d2(p, d1_grids[:,2], d2_grids[:,2])
# print(terms[0][0], terms[1][0])
# import sys
# sys.exit(0)
for i in range(n_dgrid):
for j in range(n_dgrid):
terms = ln_likelihood_at_d1d2(p, d1_grids[:,i], d2_grids[:,j])
log_d_pdf = (norm.logpdf(1000/d1_grids[:,i], plx1, plx1_err) +
norm.logpdf(1000/d2_grids[:,j], plx2, plx2_err))
ll_grid1[:,i,j] = terms[0] + log_d_pdf
ll_grid2[:,i,j] = terms[1] + log_d_pdf
l_grid1 = np.exp(ll_grid1)
lls1 = np.log([simps(simps(l_grid1[n], d2_grids[n]), d1_grids[n]) for n in range(n_data)])
l_grid2 = np.exp(ll_grid2)
lls2 = np.log([simps(simps(l_grid2[n], d2_grids[n]), d1_grids[n]) for n in range(n_data)])
return np.logaddexp(lls1, lls2), (lls1, lls2)
In [847]:
# for f in [0.1, 0.25, 0.6]:
# print(ln_likelihood_with_dist([f])[0].sum())
# ln_likelihood_with_dist([0.25])[0].sum()
In [543]:
%prun -s tottime [ln_likelihood_at_d1d2([0.25], d1_grids[:,2], d2_grids[:,2]) for i in range(10)]
In [777]:
%time ll_grid = ln_likelihood_with_dist([0.25])
In [842]:
lls = []
fs = np.linspace(0.15, 0.4, 8)
for f in fs:
lls.append(ln_likelihood_with_dist([f])[0].sum())
In [843]:
lls = np.array(lls)
plt.plot(fs, np.exp(lls - lls.max()))
plt.axvline(0.25, color='tab:green') # truth
Out[843]:
In [390]:
from astropy.io import fits
from astropy.table import Table
In [822]:
corr_cols = ['ra', 'dec', 'parallax', 'pmra', 'pmdec']
def get_tbl(icrs, plx, plx_err, pm, pm_err, rv, rv_err):
data = dict()
data['ra'] = icrs.ra.degree
data['dec'] = icrs.dec.degree
data['parallax'] = plx
data['pmra'] = pm[:,0]
data['pmdec'] = pm[:,1]
data['RV'] = rv
data['ra_error'] = np.full_like(icrs.ra.degree, 1E-10)
data['dec_error'] = np.full_like(icrs.ra.degree, 1E-10)
data['parallax_error'] = plx_err
data['pmra_error'] = pm_err[:,0]
data['pmdec_error'] = pm_err[:,1]
data['RV_err'] = rv_err
for name1 in corr_cols:
for name2 in corr_cols:
full_name = "{}_{}_corr".format(name1, name2)
data[full_name] = np.zeros_like(icrs.ra.degree)
return Table(data)
In [823]:
tbl1 = get_tbl(icrs1, plx1, plx1_err, pm1, pm_err1, rv1, rv_err1)
tbl2 = get_tbl(icrs2, plx2, plx2_err, pm2, pm_err2, rv2, rv_err2)
tbl1.write('data1.fits', overwrite=True)
tbl2.write('data2.fits', overwrite=True)
In [ ]: